import pandas as pd
import numpy as np
import warnings
from math import sqrt
warnings.filterwarnings('ignore')
from azureml.core.run import Run
from azureml.core.experiment import Experiment
from azureml.core.workspace import Workspace
from azureml.core.authentication import ServicePrincipalAuthentication
from azureml.train.automl import AutoMLConfig
import pickle
from matplotlib import pyplot as plt
from matplotlib.pyplot import figure
import mlflow
import mlflow.sklearn
from azureml.core import Workspace, Dataset
import os
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from svm import svmImp
from rf import rfImp

    

# Input subscription id, resource group, and workspace name
subscription_id = 'XXX'
resource_group = 'XXX'
workspace_name = 'XXX'
model2Test = "rf"

# 实例化workspace
workspace = Workspace(subscription_id, resource_group, workspace_name)
# Get the uri of the workspace, we will also use mlflow to record parameters and models trained.
uri = workspace.get_mlflow_tracking_uri()
mlflow.set_tracking_uri(uri)

print("uri: ", uri)

# 导入我们之前存于storage的数据
datasetTrn = Dataset.get_by_name(workspace, name='trainDataset')
print("Input dataset name: {0}; Version: {1}.".format(datasetTrn.name, datasetTrn.version))
# 转化为pandas df格式
dfTrn = datasetTrn.to_pandas_dataframe()
print("Shape of train dataset: {0}".format(dfTrn.shape))
# Features
X = dfTrn[['Temperature_C', 'Humidity', 'Wind_speed_kmph', 'Wind_bearing_degrees', 'Visibility_km', 'Pressure_millibars', 'currentWeather']].values
# Label (Values to predict)
y = dfTrn['futureWeather'].values
# 将数据差分成训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=1)
# 去均值和方差的归一化
sc = StandardScaler()
X_train = sc.fit_transform(X_train)
X_val = sc.transform(X_val)

print("Start modelling with {0} method".format(model2Test))
if model2Test == "svm":
      # 新建一个新的experiment，这个和mlflow是一个道理
      myexperiment = Experiment(workspace, "SVMTest1")
      mlflow.set_experiment("SVMTest1")
      # 在此experiment下新建一个run
      with mlflow.start_run() as new_run:
            mlflow.log_param("dataset_name", datasetTrn.name)
            mlflow.log_param("dataset_version", datasetTrn.version)
            # 实例化svmImp类
            svmImp_ = svmImp(datasetTrn, X_train, y_train, X_val, y_val,new_run)
            svmImp_.svmImpTrn()     
            svmImp_.svmImpVal()
            svmImp_.onnxModelSave()
            svmImp_.modelRegister(workspace)
            svmImp_.modelRegisterSC(sc,workspace)

      print ("run id:", new_run.info.run_id)
      mlflow.end_run()
elif model2Test == "rf":
      # 新建一个新的experiment，这个和mlflow是一个道理
      myexperiment = Experiment(workspace, "RFTest1")
      mlflow.set_experiment("RFTest1")
      # 在此experiment下新建一个run
      with mlflow.start_run() as new_run:
            mlflow.log_param("dataset_name", datasetTrn.name)
            mlflow.log_param("dataset_version", datasetTrn.version)
            # 实例化rfImp类
            rfImp_ = rfImp(datasetTrn, X_train, y_train, X_val, y_val,new_run)
            rfImp_.rfImpTrn()
            rfImp_.rfImpVal()
            rfImp_.onnxModelSave()
            rfImp_.modelRegister(workspace)
            rfImp_.modelRegisterSC(sc,workspace)
      print ("run id:", new_run.info.run_id)
      mlflow.end_run()




